package com.example.ai.embedding;

import com.example.ai.etl.MyKeywordEnricher;
import com.example.ai.etl.MyTokenTextSplitter;
import com.example.ai.etl.reader.MyJsonReader;
import com.example.ai.etl.reader.MyPagePdfDocumentReader;
import com.example.ai.etl.reader.MyTikaDocumentReader;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.JsonReader;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.writer.FileDocumentWriter;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.ClassPathResource;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;

import java.io.File;
import java.util.List;

@RestController
@RequestMapping("/verctor")
public class VectorDataInitController {
    @Value("${vector.db.path}")
    private String vectorDbPath;
    @Autowired
    VectorStore vectorStore;

    @Autowired
    MyTikaDocumentReader tikaDocumentReader;

    @Autowired
    MyTokenTextSplitter tokenTextSplitter;

    @Autowired
    BatchingStrategy tokenCountBatchingStrategy;

    @Autowired
    MyKeywordEnricher myKeywordEnricher;

    @Autowired
    MyJsonReader myJsonReader;

    @Autowired
    MyPagePdfDocumentReader myPagePdfDocumentReader;

    @GetMapping("/init/goods_json")
    public ResponseEntity<?> init() {
        initJsonData("goods.json");
        saveData();
        return ResponseEntity.ok("success");
    }

    @GetMapping("/init/job_txt")
    public ResponseEntity<?> initJobTxt() {
        initJobTxtData("job.txt");
        saveData();
        return ResponseEntity.ok("success");
    }

    @GetMapping("/init/pdf")
    public ResponseEntity<?> initPdf() {
        initPdfData("data/中华人民共和国劳动法.pdf");
        saveData();
        return ResponseEntity.ok("success");
    }

    @GetMapping("/writeToFile")
    public ResponseEntity<?> writeToFile() {
        List<Document> docs = myPagePdfDocumentReader.getDocs("data/中华人民共和国劳动法.pdf");
        FileDocumentWriter writer = new FileDocumentWriter("output.txt", true, MetadataMode.ALL, true);
        writer.accept(docs);
        return ResponseEntity.ok("success");
    }

    @GetMapping("/init/docx")
    public ResponseEntity<?> initDocx(@RequestBody EmebddingResourceReq req) {
        List<Document> documents = tikaDocumentReader.loadText(req.getClassFilePath());
        List<Document> splitedDocs = tokenTextSplitter.splitDocuments(documents);
        List<Document> enrichDocs = myKeywordEnricher.enrichDocuments(splitedDocs,req.getEnrichKeywordCount());
        this.vectorStore.add(enrichDocs);
        saveData();
        return ResponseEntity.ok("success");
    }

    @GetMapping("/enrichMetadata")
    public ResponseEntity<?> enrichMetadata(@RequestBody EnrichDocumentMetadataReq req) {
        return ResponseEntity.ok("success");
    }

    /**
     * 需要1.0.0-M7 版本才支持html Reader
     * @return
     */
    @GetMapping("/init/html")
    public ResponseEntity<?> initHtml() {
        //initHtmlData("job.txt");
        saveData();
        return ResponseEntity.ok("success");
    }

    void initJsonData(String sourceFile) {
        MyJsonMetadataGenerator myJsonMetadataGenerator = new MyJsonMetadataGenerator();
        List<Document> documents = myJsonReader.get(sourceFile,myJsonMetadataGenerator,"price","name","description","tags");
        this.vectorStore.add(documents);
    }
    void initJobTxtData(String sourceFile) {
        MyJsonMetadataGenerator myJsonMetadataGenerator = new MyJsonMetadataGenerator();
        TextReader reader = new TextReader(new ClassPathResource(sourceFile));
        reader.getCustomMetadata().put("category","job");
        List<Document> documents = reader.read();
        List<List<Document>> batches = tokenCountBatchingStrategy.batch(documents);
        for (List<Document> subBatch : batches) {
            this.vectorStore.add(subBatch);
        }
    }

    void initPdfData(String sourceFile) {
        PagePdfDocumentReader pdfReader = new PagePdfDocumentReader(new ClassPathResource(sourceFile),
                PdfDocumentReaderConfig.builder()
                        .withPageTopMargin(0)
                        .withPageExtractedTextFormatter(ExtractedTextFormatter.builder()
                                .withNumberOfTopTextLinesToDelete(0)
                                .build())
                        .withPagesPerDocument(1)
                        .build());
        List<Document> documents = pdfReader.read();
        System.out.println(documents);
    }



    void saveData(){
        ((SimpleVectorStore)this.vectorStore).save(new File(vectorDbPath));
    }
}
